#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>

#include "MSTK.h"


#ifdef __cplusplus
extern "C" {
#endif

  /* this function partitions the vertices and faces, 

  */

int MESH_Surf_Partition(Mesh_ptr mesh, Mesh_ptr *submeshes, int num) {
  int i, j, k, idx, max_nfv;
  int nv, ne, nf, nfv, nfe;
  MVertex_ptr mv, vertex, global_mv;
  MEdge_ptr me, global_me;
  MFace_ptr mf, global_mf;
  List_ptr mfverts, mfedges;

  /* basic mesh information */
  RepType rtype = MESH_RepType(mesh);
  nv = MESH_Num_Vertices(mesh);
  ne = MESH_Num_Edges(mesh);
  nf = MESH_Num_Faces(mesh);

  /* Partition the mesh */
  int *part = (int *)MSTK_malloc(nf*sizeof(int));
  int ok = MESH_PartitionWithMetis(mesh, num, &part);
  
  int part_no;
  /* global to local vertex ID mapping */
  int *v_g2l = (int *)MSTK_malloc(nv*sizeof(int));
  int index_vertex, index_face, index_edge;
  /* List of faces and vertices of this partition */
  List_ptr part_vertices, part_faces;
  /* each face belongs only to one partition, so no need to Mark  */
  int mkvid = MSTK_GetMarker();

  /* for each partition */
  for(part_no = 0; part_no < num; part_no++) {
    /* clear to 0 */
    for(j = 0; j < nv; j++)
      v_g2l[j] = 0;
    max_nfv = 0;
    part_vertices = List_New(10), part_faces = List_New(10);
    index_vertex = 0, index_face = 0, index_edge = 0;
    idx = 0, i = 0;
    while((mf = MESH_Next_Face(mesh,&idx))) {
      if(part[i++] != part_no)
	continue;
      /* add faces */
      index_face++;
      List_Add(part_faces,mf);
      
      /* add vertices */
      mfverts = MF_Vertices(mf,1,0);
      nfv = List_Num_Entries(mfverts);
      if(nfv > max_nfv)
	max_nfv = nfv;
      for(j = 0; j < nfv; j++) {
	mv = List_Entry(mfverts,j);
	if (!MEnt_IsMarked(mv,mkvid)) {
	  v_g2l[MEnt_ID(mv)-1] = index_vertex++;
	  List_Add(part_vertices,mv);
	  MEnt_Mark(mv,mkvid);
	}
      }
      List_Delete(mfverts);
    } 

    
    /* collect the coordinates and store the global vertex ID */
    int *v_l2g = (int *)MSTK_malloc(index_vertex*sizeof(int)); 
    typedef double (*COORDS)[3];
    COORDS xyz =(COORDS) MSTK_malloc(index_vertex*3*sizeof(double));
    for (i = 0; i < index_vertex; i++) {
      vertex = List_Entry(part_vertices,i);
      MV_Coords(vertex,xyz[i]);
      v_l2g[i] = MEnt_ID(vertex);
    }
    /* collect the face-vertex connectivity and store the global face ID*/ 
    int *f_l2g = (int *)MSTK_malloc(index_face*sizeof(int)); 
    int *nfvs = (int *)MSTK_malloc(index_face*sizeof(int));
    int *fvids = (int *)MSTK_malloc(index_face*max_nfv*sizeof(int));
    for (i = 0; i < index_face; i++) {
      mf = List_Entry(part_faces,i);
      mfverts = MF_Vertices(mf,1,0);
      nfv = List_Num_Entries(mfverts);
      nfvs[i] = nfv;
      f_l2g[i] = MEnt_ID(mf);
      for (j = 0; j< nfv; j++) {
	mv = List_Entry(mfverts,j);
	fvids[i*max_nfv+j] = v_g2l[MEnt_ID(mv)-1];
      }
      List_Delete(mfverts);
    }
    
    
    /* build the submesh */
    Mesh_ptr submesh = MESH_New(rtype);
    MESH_InitFromGenDesc(submesh,index_vertex,xyz,index_face,nfvs,&fvids,0,NULL,NULL,NULL,NULL);

    /* recover the vertex global ID and store the local ID into master local id*/
    for(i = 0; i < index_vertex; i++) {
      mv = MESH_Vertex(submesh,i);
      /*for mpi*/
      global_mv = MESH_VertexFromID(mesh,v_l2g[i]);
      /* copy GEnt info */
      MV_Set_GEntID(mv,MV_GEntID(global_mv));
      MV_Set_GEntDim(mv,MV_GEntDim(global_mv));
      /*if master par id is not assigned*/
      if(MV_MasterParID(global_mv) == 0) {
	/* set master partition id in the global mesh */
	MV_Set_MasterParID(global_mv,part_no+1);
	/* local is the same as global */
	MV_Set_MasterParID(mv,part_no+1);
	MV_Set_GlobalID(mv,v_l2g[i]);
      }
      /* if global is assigned, set local equal to global*/
      else {
	MV_Set_MasterParID(mv,MV_MasterParID(global_mv));
	MV_Set_GlobalID(mv,MV_GlobalID(global_mv));
      }	
      /*end for mpi*/
    }


    /* recover the face global ID and store the local ID into master local id */
    for(i = 0; i < index_face; i++) {
      mf = MESH_Face(submesh,i);
      /* for mpi */
      global_mf = MESH_FaceFromID(mesh,f_l2g[i]);
      /* copy GEnt info */
      MF_Set_GEntID(mf,MF_GEntID(global_mf));
      MF_Set_GEntDim(mf,MF_GEntDim(global_mf));
      /*set global face */
      MF_Set_MasterParID(global_mf,part_no+1);
      /*set local face*/
      MF_Set_MasterParID(mf,part_no+1);
      MF_Set_GlobalID(mf,f_l2g[i]);
      /*end for mpi*/
    }

    submeshes[part_no] = submesh;     
    
    /* clear the marker */
    List_Unmark(part_vertices,mkvid); 
    List_Delete(part_vertices);
    List_Delete(part_faces);
   
    MSTK_free(xyz);
    MSTK_free(nfvs);
    MSTK_free(fvids);
    MSTK_free(v_l2g);
    MSTK_free(f_l2g);
  }

    MSTK_FreeMarker(mkvid);
    MSTK_free(v_g2l);
    MSTK_free(part);
  
  return ok;
}
  
#ifdef __cplusplus
}
#endif

